from python_ai.common.xcommon import *
import tensorflow.compat.v1 as tf
import tensorflow as tsf
import numpy as np

tf.set_random_seed(777)

X0_batch = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 0, 1]])  #(4,3)
X1_batch = np.array([[9, 8, 7], [0, 0, 0], [6, 5, 4], [3, 2, 1]])
X2_batch = X0_batch.copy()
X3_batch = X1_batch.copy()
X4_batch = X1_batch.copy()
X5_batch = X1_batch.copy()

n_input = 3
n_hidden = 5

X0 = tf.placeholder(tf.float32, [None, n_input], 'X0')
X1 = tf.placeholder(tf.float32, [None, n_input], 'X1')
X2 = tf.placeholder(tf.float32, [None, n_input], 'X2')
X3 = tf.placeholder(tf.float32, [None, n_input], 'X3')
X4 = tf.placeholder(tf.float32, [None, n_input], 'X3')
X5 = tf.placeholder(tf.float32, [None, n_input], 'X3')

cell = tf.nn.rnn_cell.BasicRNNCell(num_units=n_hidden)

output_seqs, states = tf.nn.static_rnn(cell, [X0, X1, X2, X3, X4, X5],
                                                 dtype=tf.float32)
print(f'output_seqs: {output_seqs}')
print(f'states: {states}')

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    output_seqs_, states_ = sess.run([output_seqs, states],
                                           feed_dict={
                                               X0: X0_batch,
                                               X1: X1_batch,
                                               X2: X2_batch,
                                               X3: X3_batch,
                                               X4: X4_batch,
                                               X5: X5_batch
                                           })

for i, output in enumerate(output_seqs_):
    print(f'#{i}: {output}')
print('states:')
print(states_)
print(f'outputs: {np.shape(output_seqs_)}')
print(f'states: {np.shape(states_)}')

sep('Align to dynamic_rnn output')
output_seqs_ = np.transpose(output_seqs_, [1, 0, 2])
for i, output in enumerate(output_seqs_):
    print(f'#{i}: {output}')
print('states:')
print(states_)
print(f'outputs: {np.shape(output_seqs_)}')
print(f'states: {np.shape(states_)}')
